package com.team.merchant.websocket;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.team.merchant.domain.vo.StatsMessage;
import com.team.merchant.service.StatsService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 商户WebSocket处理器
 * 处理微信小程序的WebSocket连接
 */
@Component
@RequiredArgsConstructor
@Slf4j
public class MerchantWebSocketHandler extends TextWebSocketHandler {

    private final ObjectMapper objectMapper;
    private final StatsService statsService;
    
    // 存储订阅了各个主题的会话
    private final Map<String, ConcurrentHashMap<String, WebSocketSession>> topicSubscribers = new ConcurrentHashMap<>();
    
    // 存储认证过的会话
    private final Map<String, String> authenticatedSessions = new ConcurrentHashMap<>();
    
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("新的WebSocket连接已建立: {}", session.getId());
        
        // 从URL参数中获取token
        String token = extractTokenFromSession(session);
        if (token != null && !token.isEmpty()) {
            // 记录会话已认证
            authenticatedSessions.put(session.getId(), token);
            log.info("会话已通过URL参数认证: {}", session.getId());
        }
    }
    
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        log.info("收到WebSocket消息: {}", payload);
        
        try {
            JsonNode jsonNode = objectMapper.readTree(payload);
            String type = jsonNode.has("type") ? jsonNode.get("type").asText() : "";
            
            switch (type) {
                case "auth":
                    handleAuth(session, jsonNode);
                    break;
                case "subscribe":
                    handleSubscribe(session, jsonNode);
                    break;
                default:
                    log.warn("未知的消息类型: {}", type);
                    sendErrorMessage(session, "未知的消息类型: " + type);
            }
        } catch (Exception e) {
            log.error("处理WebSocket消息失败: {}", e.getMessage(), e);
            sendErrorMessage(session, "处理消息失败: " + e.getMessage());
        }
    }
    
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("WebSocket连接已关闭: {}, 状态: {}", session.getId(), status);
        
        // 清理认证信息
        authenticatedSessions.remove(session.getId());
        
        // 清理订阅信息
        for (Map.Entry<String, ConcurrentHashMap<String, WebSocketSession>> entry : topicSubscribers.entrySet()) {
            entry.getValue().remove(session.getId());
        }
    }
    
    /**
     * 处理认证消息
     */
    private void handleAuth(WebSocketSession session, JsonNode jsonNode) {
        try {
            if (jsonNode.has("token")) {
                String token = jsonNode.get("token").asText();
                // 这里可以添加token验证逻辑
                
                // 记录会话已认证
                authenticatedSessions.put(session.getId(), token);
                log.info("会话已通过消息认证: {}", session.getId());
                
                // 发送认证成功消息
                sendMessage(session, "{\"type\":\"auth\",\"status\":\"success\"}");
            } else {
                log.warn("认证消息缺少token字段");
                sendErrorMessage(session, "认证消息缺少token字段");
            }
        } catch (Exception e) {
            log.error("处理认证消息失败: {}", e.getMessage(), e);
            sendErrorMessage(session, "处理认证消息失败: " + e.getMessage());
        }
    }
    
    /**
     * 处理订阅消息
     */
    private void handleSubscribe(WebSocketSession session, JsonNode jsonNode) {
        try {
            // 检查是否已认证
            if (!authenticatedSessions.containsKey(session.getId())) {
                log.warn("未认证的会话尝试订阅: {}", session.getId());
                sendErrorMessage(session, "请先进行认证");
                return;
            }
            
            if (jsonNode.has("topic")) {
                String topic = jsonNode.get("topic").asText();
                
                // 添加到订阅列表
                topicSubscribers.computeIfAbsent(topic, k -> new ConcurrentHashMap<>())
                        .put(session.getId(), session);
                
                log.info("会话 {} 已订阅主题: {}", session.getId(), topic);
                
                // 发送订阅成功消息
                sendMessage(session, "{\"type\":\"subscribe\",\"status\":\"success\",\"topic\":\"" + topic + "\"}");
                
                // 立即发送一次最新数据
                if ("/topic/stats".equals(topic)) {
                    sendStatsMessage(session);
                }
            } else {
                log.warn("订阅消息缺少topic字段");
                sendErrorMessage(session, "订阅消息缺少topic字段");
            }
        } catch (Exception e) {
            log.error("处理订阅消息失败: {}", e.getMessage(), e);
            sendErrorMessage(session, "处理订阅消息失败: " + e.getMessage());
        }
    }
    
    /**
     * 发送统计数据消息
     */
    private void sendStatsMessage(WebSocketSession session) {
        try {
            StatsMessage statsMessage = statsService.getLatestStats();
            String message = objectMapper.writeValueAsString(statsMessage);
            sendMessage(session, message);
        } catch (Exception e) {
            log.error("发送统计数据失败: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 广播统计数据到所有订阅者
     */
    public void broadcastStats(StatsMessage statsMessage) {
        try {
            String message = objectMapper.writeValueAsString(statsMessage);
            broadcast("/topic/stats", message);
        } catch (Exception e) {
            log.error("广播统计数据失败: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 向指定主题的所有订阅者广播消息
     */
    public void broadcast(String topic, String message) {
        ConcurrentHashMap<String, WebSocketSession> subscribers = topicSubscribers.get(topic);
        if (subscribers != null && !subscribers.isEmpty()) {
            log.info("向主题 {} 的 {} 个订阅者广播消息", topic, subscribers.size());
            
            for (WebSocketSession session : subscribers.values()) {
                try {
                    if (session.isOpen()) {
                        sendMessage(session, message);
                    }
                } catch (Exception e) {
                    log.error("向会话 {} 发送消息失败: {}", session.getId(), e.getMessage());
                }
            }
        } else {
            log.info("主题 {} 没有订阅者", topic);
        }
    }
    
    /**
     * 发送错误消息
     */
    private void sendErrorMessage(WebSocketSession session, String errorMessage) {
        try {
            String message = "{\"type\":\"error\",\"message\":\"" + errorMessage + "\"}";
            sendMessage(session, message);
        } catch (Exception e) {
            log.error("发送错误消息失败: {}", e.getMessage(), e);
        }
    }
    
    /**
     * 发送WebSocket消息
     */
    private void sendMessage(WebSocketSession session, String message) throws IOException {
        if (session.isOpen()) {
            session.sendMessage(new TextMessage(message));
        }
    }
    
    /**
     * 从会话中提取token
     */
    private String extractTokenFromSession(WebSocketSession session) {
        try {
            String query = session.getUri().getQuery();
            if (query != null && query.contains("token=")) {
                String[] params = query.split("&");
                for (String param : params) {
                    if (param.startsWith("token=")) {
                        return param.substring(6); // "token=".length() == 6
                    }
                }
            }
        } catch (Exception e) {
            log.error("从会话中提取token失败: {}", e.getMessage(), e);
        }
        return null;
    }
} 